import os
import requests
import json
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class AIService:
    """
    AI助手服务类，用于与阿里云百炼平台通信
    """
    def __init__(self):
        self.api_key = os.getenv('DASHSCOPE_API_KEY')
        self.model = os.getenv('DASHSCOPE_MODEL')
        self.api_url = os.getenv('DASHSCOPE_API_URL')
        
        if not self.api_key or not self.model or not self.api_url:
            raise ValueError("阿里云百炼平台配置缺失，请检查.env文件")
    
    def generate_response(self, prompt, history=None):
        """
        生成AI回复
        
        Args:
            prompt (str): 用户输入的问题
            history (list, optional): 对话历史记录，格式为 [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
            
        Returns:
            dict: 包含AI回复的字典
        """
        if history is None:
            history = []
            
        # 构建请求头
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        # 构建请求体
        messages = history.copy()
        messages.append({"role": "user", "content": prompt})
        
        payload = {
            "model": self.model,
            "input": {
                "messages": messages
            },
            "parameters": {
                "temperature": 0.7,
                "top_p": 0.8,
                "result_format": "message"
            }
        }
        
        try:
            # 发送请求
            response = requests.post(
                self.api_url or "",  # 添加默认值防止None
                headers=headers,
                data=json.dumps(payload)
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            # 提取AI回复
            ai_response = result.get('output', {}).get('choices', [{}])[0].get('message', {}).get('content', '')
            
            return {
                "success": True,
                "message": ai_response,
                "raw_response": result
            }
            
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "error": f"请求失败: {str(e)}"
            }
        except (KeyError, IndexError, ValueError) as e:
            return {
                "success": False,
                "error": f"解析响应失败: {str(e)}"
            }
    
    def generate_stream_response(self, prompt, history=None):
        """
        生成AI流式回复
        
        Args:
            prompt (str): 用户输入的问题
            history (list, optional): 对话历史记录
            
        Yields:
            str: AI回复的文本片段
        """
        if history is None:
            history = []
            
        # 构建请求头
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",  # 启用流式响应
            "X-DashScope-SSE": "enable"     # 阿里云百炼平台流式输出标志
        }
        
        # 构建请求体
        messages = history.copy()
        messages.append({"role": "user", "content": prompt})
        
        payload = {
            "model": self.model,
            "input": {
                "messages": messages
            },
            "parameters": {
                "temperature": 0.7,
                "top_p": 0.8,
                "result_format": "message",
                "incremental_output": True  # 启用增量输出
            }
        }
        
        try:
            # 发送流式请求
            response = requests.post(
                self.api_url or "",  # 添加默认值防止None
                headers=headers,
                data=json.dumps(payload),
                stream=True  # 启用流式传输
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 处理流式响应
            for line in response.iter_lines(decode_unicode=True):
                if line:
                    # 解析SSE事件
                    if line.startswith("data:"):
                        data = line[5:].strip()  # 移除"data:"前缀
                        if data:
                            try:
                                # 解析JSON数据
                                result = json.loads(data)
                                # 提取内容
                                content = result.get('output', {}).get('choices', [{}])[0].get('message', {}).get('content', '')
                                yield content
                            except json.JSONDecodeError:
                                # 如果不是JSON格式，直接返回原始数据
                                yield data
                                
        except requests.exceptions.RequestException as e:
            yield f"错误: 请求失败: {str(e)}"
        except Exception as e:
            yield f"错误: {str(e)}"

# 创建AI服务实例
ai_service = AIService()